import os
import bpy
import bmesh
import math
import numpy as np
from mathutils import Euler, Matrix, Vector


def look_at(location):
    """
    Get object pose from a viewpoint location
    # Taken from https://github.com/thodan/bop_toolkit/blob/master/bop_toolkit_lib/view_sampler.py#L216
    IMPORTANT: output of this function is the object pose defined in OPENGL coordinate convention
    """
    f = -np.array(location)  # Forward direction.
    f /= np.linalg.norm(f)

    u = np.array([0.0, 0.0, 1.0])  # Up direction.
    s = np.cross(f, u)  # Side direction.
    if np.count_nonzero(s) == 0:
        # f and u are parallel, i.e. we are looking along or against Z axis.
        s = np.array([1.0, 0.0, 0.0])
    s /= np.linalg.norm(s)
    u = np.cross(s, f)  # Recompute up.
    R = np.array([[s[0], s[1], s[2]],
                  [u[0], u[1], u[2]],
                  [-f[0], -f[1], -f[2]]])
    t = - R.dot(np.array(location).reshape((3, 1)))
    obj_pose = np.zeros((4, 4))
    obj_pose[:3, :3] = R
    obj_pose[:3, 3] = -t.reshape(-1)
    obj_pose[3, 3] = 1
    return obj_pose

def opencv2opengl(cam_matrix_world):
    """
    Change coordinate system from OpenCV to OpenGL or from OpenGL to OpenCV
    """
    from scipy.spatial.transform import Rotation as R
    rot180x = R.from_euler('x', 180, degrees=True).as_matrix()
    rotation = cam_matrix_world[:3, :3]
    translation = cam_matrix_world[:3, 3]
    output = np.copy(cam_matrix_world)
    output[:3, :3] = np.asarray(Matrix(rot180x) @ Matrix(rotation).to_3x3())
    output[:3, 3] = np.asarray(Matrix(rot180x) @ Vector(translation))
    return output


def get_pose():
    camera=bpy.data.objects["Camera"]
    obj=bpy.data.objects['Icosphere']
    obj_world_matrix=obj.matrix_world
    camera_world_matrix=camera.matrix_world
    obj_to_camera_matrix = np.dot(camera_world_matrix.inverted(), obj_world_matrix)
    return obj_to_camera_matrix



def get_camera_positions(nSubDiv):
    """
    * Construct an icosphere
    * subdived
    """
    os.makedirs(r"E:\graguate\blender\datasets1-sub4\matrix",exist_ok=True)
    os.makedirs(r"E:\graguate\blender\datasets1-sub4\position",exist_ok=True)
    os.makedirs(r"E:\graguate\blender\datasets1-sub4\angles",exist_ok=True)

    bpy.ops.mesh.primitive_ico_sphere_add(location=(0, 0, 0), enter_editmode=True)
    icos = bpy.context.object
    me = icos.data

    # -- subdivide and move new vertices out to the surface of the sphere
    #    nSubDiv = 3
    for i in range(nSubDiv):
        bpy.ops.mesh.subdivide()

        bm = bmesh.from_edit_mesh(me)
        for v in bm.verts:
            l = math.sqrt(v.co[0] ** 2 + v.co[1] ** 2 + v.co[2] ** 2)
            v.co[0] /= l
            v.co[1] /= l
            v.co[2] /= l
        bmesh.update_edit_mesh(me)

    # -- cut away zero elevation
    bm = bmesh.from_edit_mesh(me)
    sel = [v for v in bm.verts if v.co[2] == 0]
    bmesh.ops.delete(bm, geom=sel, context="FACES")
    bmesh.update_edit_mesh(me)

    # convert vertex positions to az,el
    positions = []
    angles = []
    obj_poses=[]
    bm = bmesh.from_edit_mesh(me)
    for idx,v in enumerate(bm.verts):
        x = v.co[0]
        y = v.co[1]
        z = v.co[2]
        az = math.atan2(x, y)  # *180./math.pi
        el = math.atan2(z, math.sqrt(x ** 2 + y ** 2))  # *180./math.pi
        # positions.append((az,el))
        angles.append((el, az)) #极坐标
        positions.append((x, y, z))

        ''' 获取物体转换位姿'''

        ''' 获取物体的位姿 '''
        obj_pose=get_pose()
        obj_poses.append(obj_pose)

        #保存位姿
        np.savetxt(os.path.join(r"E:\graguate\blender\datasets1-sub4\matrix","{}.txt".format(idx)),obj_pose,fmt="%.6f")
        np.savetxt(os.path.join(r"E:\graguate\blender\datasets1-sub4\position","{}.txt".format(idx)),(x, y, z),fmt="%.6f")
        np.savetxt(os.path.join(r"E:\graguate\blender\datasets1-sub4\angles","{}.txt".format(idx)),(el, az),fmt="%.6f")
    
    bpy.ops.object.editmode_toggle()
    return angles, positions,obj_poses

'''
gpu 渲染

'''

'''
print("start............")
level =4
angles,position_icosphere,poses = get_camera_positions(level)
angles = np.array(angles)
position_icosphere = np.array(position_icosphere)
poses=np.array(poses)
print("end............")
'''
